import gradio as gr
import pickle
import numpy as np
from PIL import Image
import cv2
from sklearn.datasets import load_digits

# 1. 加载模型并验证
try:
    with open('best_knn_model.pkl', 'rb') as f:
        knn_model = pickle.load(f)
    # 验证模型是否正常（测试训练集样本）
    digits = load_digits()
    test_samples = digits.data[0:5]
    test_labels = digits.target[0:5]
    predictions = knn_model.predict(test_samples)
    print("=== 模型验证结果 ===")
    for i in range(5):
        print(f"样本{i+1}：真实标签{test_labels[i]} → 模型预测{predictions[i]}")
    if len(set(predictions)) == 1:
        print("警告：模型可能异常（预测结果单一），建议重新运行 optimal_knn.py！")
except FileNotFoundError:
    print("错误：未找到 best_knn_model.pkl，请先运行 optimal_knn.py 生成模型！")
    exit()

# 2. 修复后的预测函数
def predict_digit(input_data):
    if input_data is None:
        return "请先在画板上绘制一个数字（用白色粗笔）"
    
    # 提取图像数据
    if isinstance(input_data, dict):
        input_data = input_data.get('composite', input_data.get('layers', [None])[-1])
    if input_data is None:
        return "未检测到绘制内容，请重新绘制"
    
    # 处理图像格式
    if not isinstance(input_data, np.ndarray):
        try:
            input_data = np.array(input_data)
        except:
            return "输入格式错误，请重新绘制"
    # 处理RGBA转RGB
    if input_data.shape[2] == 4:
        input_data = input_data[:, :, :3]
    if len(input_data.shape) != 3 or input_data.shape[2] != 3:
        return f"图像格式错误（需3通道RGB），当前: {input_data.shape}"
    
    # 核心预处理（对齐训练数据）
    try:
        # 灰度化 + 低阈值二值化（捕捉淡色笔画）
        gray_img = Image.fromarray(input_data.astype('uint8'), 'RGB').convert('L')
        img_array = np.array(gray_img)
        _, binary_img = cv2.threshold(img_array, 80, 255, cv2.THRESH_BINARY)
        
        # 缩放为8x8 + 颜色反转 + 像素范围对齐（替换ANTIALIAS为LANCZOS）
        resized_img = Image.fromarray(binary_img).resize((8, 8), Image.LANCZOS)
        resized_array = np.array(resized_img)
        processed_array = (255.0 - resized_array) * (16 / 255.0)  # 匹配训练数据的0-16像素范围
        processed_array[processed_array < 1] = 0  # 清除背景噪点
        
        # 预测
        flattened_array = processed_array.flatten()
        prediction = knn_model.predict([flattened_array])
        return f"预测结果: {str(prediction[0])}"
    except Exception as e:
        return f"处理出错: {str(e)}"

# 3. 启动服务（指定端口，避免冲突）
with gr.Blocks(title="手写数字识别") as iface:
    gr.Markdown("# 手写数字识别")
    gr.Markdown("绘制要求：1. 用白色粗笔 2. 占满画板 3. 无多余线条")
    
    with gr.Row():
        sketchpad = gr.Sketchpad(label="绘制区域", height=350)
        output = gr.Label(label="识别结果", show_label=True)
    
    with gr.Row():
        submit_btn = gr.Button("提交", variant="primary")
        clear_btn = gr.Button("清除")
    
    # 绑定提交按钮的点击事件
    submit_btn.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=output
    )
    
    # 绑定清除按钮的点击事件（清空画板）
    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=sketchpad
    )

if __name__ == "__main__":
    print("正在启动服务... 访问地址：http://127.0.0.1:7862")
    iface.launch(share=False, server_port=7862, server_name="0.0.0.0")